Skip to content

Fix boundary coverage and top-k packing in three Triton tasks - #120

Draft
irvineoy wants to merge 2 commits into
mainfrom
fix/triton-boundary-correctness
Draft

irvineoy wants to merge 2 commits into
mainfrom
fix/triton-boundary-correctness

Conversation

@irvineoy

@irvineoy irvineoy commented Oct 8, 2026 •

Copy link
Copy Markdown
Collaborator

The token-combine and temperature controls missed several layout and boundary cases, and bitmatrix packing dropped assignments after the first 32 top-k slots. This change fixes packing across all tiles and strengthens the three tasks' correctness and timing checks.

  • Cover zero/one/two draft tokens, request remapping, padded draft rows, and read-only input/padding checks in token combination.
  • Cover remapped temperature requests, padded logits, and a partial final block at vocabulary width 8449. A skip-tail implementation is rejected by the native control.
  • Enforce the documented three-argument apply_temperature(logits, idx_mapping, temperature) signature, rejecting optional extras, defaults, variadic arguments, and changed parameter names or kinds before the harness runs.
  • Test top-k 33/65 with five distinct rows and 70 experts. Missing-middle-tile and wrong-row implementations fail the native gate. Add three correctness-only strided-input controls; the wrapper makes those views contiguous before calling the unchanged Triton kernel. Include the boundary contract in the task prompt and distinguish imported controls from repository-authored additions.
  • Check every packing output charged to a captured graph replay, require outputs that do not share bytes, then poison and recheck every output on perturbed-input replay. This rejects shared-output and stale-intermediate-output shortcuts. Disjoint views of one allocation, including interleaved strided views, remain valid.

The original five scored workloads per task, tolerances, warmups, sample counts, and canonical timing helper are unchanged. Added workloads are correctness-only. The packing kernel retains its original single-tile path for the scored shapes; the wrapper's layout normalization is a no-op for those contiguous inputs.

Validation: all three tasks have fresh framework-finalized PASS reports on AMD Instinct MI355X (gfx950), with all 12 checks passing per task. The same job passed all 52 focused GPU tests with no skips. CPU regression: 2361 passed, 10 GPU-only tests skipped locally. make check-perf-helpers passed. Environment: Linux-6.8.0-107-generic-x86_64-with-glibc2.39, Python 3.12.3, PyTorch 2.11.0+rocm10.0.0, HIP 7.15.26333, Triton 3.8.0; Validator: task_validator. Docker: lmsysorg/sglang@sha256:e20849665c105d389ef91d23c0dc73931aaa6f02056dd10e7b43e4f16c79df69.

Scored case Previous harness (µs) Updated harness (µs) Change
perf1 7.220 7.196 -0.33%
perf2 7.281 7.209 -0.99%
perf3 7.061 7.195 +1.91%
perf4 11.828 11.747 -0.69%
perf5 12.133 12.220 +0.72%

The paired comparison uses the same packing Triton kernel AST before and after the harness and wrapper changes, with three alternating AB/BA/AB pairs, native graph timing, 10 warmups, and 100 samples. The historical package passes its ten numeric cases and the updated package passes thirteen; the older harness still has the demonstrated validation gap. These measurements check the effect of the harness and layout-normalizing wrapper together; they are not a kernel speedup claim. Qualification is limited to MI355X and the pinned runtime. The checks reject the demonstrated shared-output bypass, but do not prove independent computation for every possible cache-and-copy implementation on identical inputs.

CPU checks:

python3 -m pytest -q tests/test_triton_task_migration_v2.py tests/test_triton_token_replay_v2.py tests/test_combine_decode_controls.py tests/test_temperature_layout_controls.py tests/test_pack_bitmatrix_topk_boundary.py tests/test_pack_bitmatrix_benchmark_integrity.py tests/test_pr107_main_integration.py tests/test_task_prompt_v2.py
make check-perf-helpers

GPU checks inside the pinned runtime:

python3 -m pytest -q tests/test_combine_decode_controls.py tests/test_temperature_layout_controls.py tests/test_pack_bitmatrix_topk_boundary.py tests/test_pack_bitmatrix_benchmark_integrity.py

The serial Docker validator selects triton2triton/vllm/triton_combine_sampled_and_draft_tokens, triton2triton/vllm/triton_temperature, and triton2triton/vllm/triton_pack_bitmatrix through the normal main.py path. Source and control manifests are checked before and after execution. Each finalized report is matched to the current task sources, initial source inventory, and materialized baseline.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant